# 矢量解码器（Transformer/GPT）

from transformers import GPT2LMHeadModel, GPT2Config
import torch.nn as nn

class GlyphDecoder(GPT2LMHeadModel):
    def __init__(self, struct_dim=256):
        # 修正维度配置（确保256可被num_attention_heads整除）
        config = GPT2Config(
            n_embd=256,
            n_head=8,          # 修改为8（256/8=32）
            n_layer=4,
            n_positions=1024,
            vocab_size=1,       # 伪参数
            n_ctx=1024,
            bos_token_id=0,     # 新增必要参数
            eos_token_id=1      # 新增必要参数
        )
        super().__init__(config)
        
        # 增加结构特征投影层
        self.struct_proj = nn.Linear(struct_dim, config.n_embd)
        
    def forward(self, style_feat):
        # [batch, seq=1, hidden_size]
        inputs_embeds = self.struct_proj(style_feat).unsqueeze(1)
        return super().forward(inputs_embeds=inputs_embeds).last_hidden_state
        # 调整输入维度以匹配风格+结构特征
        self.transformer.wte = nn.Linear(512, config.n_embd)